/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file
    Contains monochromatic and polychromatic policy classes for
    scatterers and dust grains used by the scatterer and dust_model
    classes. */

// $Id$

#ifndef __chromatic_policy__
#define __chromatic_policy__


#include <vector> 
#include "boost/shared_ptr.hpp"

#include "mcrx-types.h"
#include "blitz/array.h"
#include "blitz/ops.h"
#include "ray.h"

namespace mcrx {
  class generic_biaser;
  class monochromatic_biaser;
  class polychromatic_biaser;
  template <typename> class generic_chromatic_policy;
  class monochromatic_policy;
  class polychromatic_policy;
  template <typename> class monochromatic_scatterer_policy;
  template <typename> class monochromatic_dust_model_policy;
  template <typename> class polychromatic_scatterer_policy;
  template <typename> class polychromatic_dust_model_policy;
}

/**** bias */

/** The generic biaser is used by the generic chromatic policy. It
    does nothing, since in general we don't know what to do... */
class mcrx::generic_biaser {
  bool operator==(const generic_biaser&) {return true;};
  bool operator!=(const generic_biaser&) {return false;};
};

/** This class implements the biasing policy for monochromatic
    radiative transfer, in which case the functions are trivial.  */
class mcrx::monochromatic_biaser {
public:
  bool operator==(const monochromatic_biaser&) {return true;};
  bool operator!=(const monochromatic_biaser&) {return false;};

  /// The reference value is the value itself.
  T_float reference1(T_float v) const {return v;};

  /// The array reference is also the array itself.
  const array_1& reference2(const array_1& v) const {return v;};
  template<typename T>
  blitz::_bz_ArrayExpr<T> reference2(blitz::_bz_ArrayExpr<T> v) const {
    return v;};

  /// Component zero is just the 0-component of the array itself
  T_float reference2_comp0(const array_1& v) const {return v(0);};

  /// Calculate the bias factor for an array. In monochromatic RT, the
  /// bias factor is always one.
  T_float bias_factor(T_float v) const {return 1;};
};


/** This class implements the biasing policy for polychromatic
    radiative transfer.  The object knows which index is the reference
    and can thus extract the reference value of an array or calculate
    the bias factor.  */
class mcrx::polychromatic_biaser {
private:
  /// The reference index into the arrays.
  int refind_;
public:
  // we need a default constructor so we just make it zero
  polychromatic_biaser() : refind_(0) {};
  /// Constructor takes the reference index value.
  polychromatic_biaser(int r) : refind_(r) {};

  bool operator==(const polychromatic_biaser& rhs) {
    return refind_==rhs.refind_;};
  bool operator!=(const polychromatic_biaser& rhs) {
    return refind_!=rhs.refind_;};
  int reference_index() const {return refind_;};

  /// Extract the reference value from an array.
  T_float reference1(const array_1& v) const {return v(refind_);};
  template<typename T>
  T_float reference1(blitz::_bz_ArrayExpr<T> v) const {return v(refind_);};

  /// Extract the reference column from a 2D array.
  array_1 reference2(const array_2& v) const {
    return v(blitz::Range::all(), refind_);};

  /// Extract the reference column of the first row of a 2D
  /// array. (Used in dust_model.h for special code path if we only
  /// have one density.)
  T_float reference2_comp0(const array_2& v) const {
    return v(0, refind_);};

  /** Calculate the bias factor for an array, returned as an
      expression template. The bias factor is the array divided by the
      reference value. */
  blitz::_bz_ArrayExpr<blitz::_bz_ArrayWhere<blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprBinaryOp<blitz::_bz_ArrayExpr<blitz::FastArrayIterator<double, 1> >, blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprConstant<double> >, blitz::Equal<double, double> > >, blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprConstant<double> >, blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprBinaryOp<blitz::_bz_ArrayExpr<blitz::FastArrayIterator<double, 1> >, blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprConstant<double> >, blitz::Divide<double, double> > > > >
  //  int
  //blitz::_bz_ArrayExpr<blitz::_bz_ArrayWhere<blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprBinaryOp<blitz::_bz_ArrayExpr<blitz::FastArrayIterator<mcrx::T_float={double}, 1>>, blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprConstant<mcrx::T_float={double}>>, blitz::Equal<mcrx::T_float={double}, mcrx::T_float={double}>>>, blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprConstant<mcrx::T_float={double}>>,        blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprBinaryOp<blitz::_bz_ArrayExpr<blitz::FastArrayIterator<mcrx::T_float={double}, 1>>, blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprConstant<mcrx::T_float={double}>>, blitz::Divide<mcrx::T_float={double}, mcrx::T_float={double}>>>>>

  bias_factor(const array_1& v) const {
    // careful we're not returning an ET which refers to an
    // out-of-scope array here? i think its ok because ref1(v) is a
    // float.  should take care of special fact that 0/0 should be
    // treated as 1 because is rare cases the scattering_tau can
    // underflow to zero
    return where(v==reference1(v), 1.0, v/reference1(v));};

  template<typename T>
  array_1
  /* can't return an expr here because it will contain a reference to
     out-of-scope
     typename blitz::BzBinaryExprResult<blitz::Divide,
				     array_1,
				     T_float>::T_result*/
  bias_factor(blitz::_bz_ArrayExpr<T> vv) const {
    const array_1 v(vv);
    // this can't be done with ETs directly, because
    // the () operator is not defined for ArrayExprs. (It may have been added.)
    //return bias_factor(v);
    array_1 ret(where(v==reference1(v), 1.0, v/reference1(v)));
    ret.threadLocal();
    return ret;
  };
};

/**** basic chromatic policy */

/** This class defines the types used for "generic" transfer problems
    which in practice means projecting sources of emission directly
    onto the cameras, without any scattering or absorption. */
template <typename content_type>
class mcrx::generic_chromatic_policy {
public:
  /// The type of wavelengths, which is a template parameter. It's
  /// supposed to be only a container.
  typedef content_type T_lambda;
  /// The type of the monochromatic biasing policy.
  typedef generic_biaser T_biaser;
  /// The type of the base policy is itself, since this class is a
  /// dummy for all the chromatic_policy classes.
  typedef generic_chromatic_policy<T_lambda> T_base_policy;
  /// The scatterer type is a dummy.
  typedef void T_scatterer;

};

/** This class defines the generic (non-template) types used for
    monochromatic radiative transfer. */
class mcrx::monochromatic_policy {
public:
  /// The type of wavelengths. In monochromatic RT, it's just a float.
  typedef T_float T_lambda;
  /// The type of the monochromatic biasing policy.
  typedef monochromatic_biaser T_biaser;
};

/** This class defines the generic (non-template) types used for
    polychromatic radiative transfer. */
class mcrx::polychromatic_policy {
public:
  /// The type of wavelengths. In polychromatic RT, it's a 1D array.
  typedef array_1 T_lambda;
  /// The type of the polychromatic biasing policy.
  typedef polychromatic_biaser T_biaser;
};

/**** scatterer_policy */

/** This class implements the chromatic policy for monochromatic
    radiative transfer used by the scatterer class. It defines the
    types used in monochromatic calculations. */
template <typename scatterer_type> 
class mcrx::monochromatic_scatterer_policy : public mcrx::monochromatic_policy {
public:
  /// The type of the base policy.
  typedef monochromatic_policy T_base_policy;
  /// The type of the scatterer, which is a template parameter.
  typedef scatterer_type T_scatterer;
  /// The type of a vector holding pointers to scatterers.
  typedef std::vector<boost::shared_ptr<T_scatterer> > T_scatterer_vector;
  /// The type of the associated dust model policy.
  typedef monochromatic_dust_model_policy<T_scatterer> T_dust_model_policy;
};


/** This class implements the chromatic policy for polychromatic
    radiative transfer used by the scatterer class. It defines the
    types used in polychromatic calculations. */
template <typename scatterer_type> 
class mcrx::polychromatic_scatterer_policy : public mcrx::polychromatic_policy {
public:
  /// The type of the base policy.
  typedef polychromatic_policy T_base_policy;
  /// The type of the scatterer, which is a template parameter.
  typedef scatterer_type T_scatterer;
  /// The type of a vector holding pointers to scatterers.
  typedef std::vector<boost::shared_ptr<T_scatterer> > T_scatterer_vector;
  /// The type of the associated dust model policy.
  typedef polychromatic_dust_model_policy<T_scatterer> T_dust_model_policy;
};

/**** dust_model_policy */

/** This class implements the chromatic policy for monochromatic
    radiative transfer used by the scatterer class. It inherits the
    scatterer policy, adds some type definitions, and implements some
    functions that are different in mono- and polychromatic
    calculations. */
template <typename scatterer_type> 
class mcrx::monochromatic_dust_model_policy :
  public mcrx::monochromatic_scatterer_policy<scatterer_type> {
public:
  /// The type of the base policy.
  typedef monochromatic_policy T_base_policy;
  /// The type of the scatterer policy.
  typedef typename mcrx::monochromatic_scatterer_policy<scatterer_type> T_scatterer_policy;
  /// Forward definition from base class.
  typedef typename T_scatterer_policy::T_lambda T_lambda;
  /// Forward definition from base class.
  typedef typename T_scatterer_policy::T_scatterer_vector T_scatterer_vector;
  /// Forward definition from base class.
  typedef typename T_scatterer_policy::T_biaser T_biaser;
  /// The type of an array of opacities, which is a 1D array in
  /// monochromatic calculations. The array is indexed by (species).
  typedef array_1 T_opacity_array;
  /// The return type of the reference_abs_lengths function.
  typedef blitz::BzBinaryExprResult<blitz::Multiply,
				    T_densities,
				    T_opacity_array>::T_result 
  RT_reference_abs_lengths;

  static T_lambda total_abs_length (const T_opacity_array&);

  template<typename T>
  static T_lambda total_abs_length (blitz::_bz_ArrayExpr<T> expr);

  static T_float total_reference_abs_length (const T_densities&, 
					     const T_opacity_array&, T_biaser);

  static
  blitz::BzBinaryExprResult<blitz::Multiply,
    T_densities,
    T_opacity_array>::T_result
  abs_length_array (const T_densities&,
		    const T_opacity_array&);
  static T_densities reference_abs_lengths (const T_densities&, const T_opacity_array&, T_biaser);
  static T_lambda extract_scatterer_component (const T_opacity_array&, int);

  static void set_wavelength_prefix (T_opacity_array& kappa,
				     T_scatterer_vector& scatterers,
				     const T_lambda& lambda);
  static void set_wavelength_postfix (T_opacity_array& kappa,
				      T_scatterer_vector& scatterers,
				      const T_lambda& lambda);
};


/** This class implements the chromatic policy for polychromatic
    radiative transfer used by the scatterer class. It inherits the
    scatterer policy, adds some type definitions, and implements some
    functions that are different in mono- and polychromatic
    calculations. */
template <typename scatterer_type> 
class mcrx::polychromatic_dust_model_policy :
  public mcrx::polychromatic_scatterer_policy<scatterer_type> {
public:
  /// The type of the base policy.
  typedef polychromatic_policy T_base_policy;
  /// The type of the scatterer policy.
  typedef typename mcrx::polychromatic_scatterer_policy<scatterer_type> T_scatterer_policy;
  /// Forward definition from base class.
  typedef typename T_scatterer_policy::T_lambda T_lambda;
  /// Forward definition from base class.
  typedef typename T_scatterer_policy::T_scatterer_vector T_scatterer_vector;
  /// Forward definition from base class.
  typedef typename T_scatterer_policy::T_biaser T_biaser;
  /// The type of an array of opacities, which is a 2D array in
  /// polychromatic calculations. The array is indexed by (species, lambda).
  typedef array_2 T_opacity_array; 
  /// The return type of the reference_abs_lengths function.
  typedef blitz::BzBinaryExprResult<blitz::Multiply,
				    T_densities,
				    array_1>::T_result 
  RT_reference_abs_lengths;


  /** Calculate the total absorption length (summed over all species) by
  summing the absorption length array, returned as as expression template. */
  static 

  typename blitz::BzReductionResult<
    blitz::ReduceSum, 
    1, 
    typename blitz::BzIndexmapResult<
      T_opacity_array,
      1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
      >::T_result
    >::T_result

			     
  //blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprReduce<blitz::_bz_ArrayExpr<blitz::ArrayIndexMapping<blitz::FastArrayIterator<double, 2>, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0> >,  1, blitz::ReduceSum<double, double> > >
  total_abs_length (const T_opacity_array& abslen)
  {
    using namespace blitz;
    return sum (abslen(tensor::j, tensor::i), tensor::j);
  }
  

  /** Traits class containing the return type definition for
      2-argument version of total_abs_length. */
  template <typename T> struct RT_total_abs_length2 {
    typedef typename blitz::BzReductionResult<
      blitz::ReduceSum, 1, 
      typename blitz::BzBinaryExprResult<
	blitz::Multiply,
	T,
	typename blitz::BzIndexmapResult<
	  T_opacity_array, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0>::T_result
	>::T_result
      >::T_result T_result;
    //typedef blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprReduce<blitz::_bz_ArrayExpr<blitz::_bz_ArrayExprBinaryOp<typename blitz::asExpr<T>::T_expr, blitz::_bz_ArrayExpr<blitz::ArrayIndexMapping<blitz::FastArrayIterator<double, 2>, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0> >, blitz::Multiply<double, double> > >, 1, blitz::ReduceSum<double, double> > > T_result;
  };

  /** Computes the total absorption length (mean free path) by summing
      over all density species. This is called from dust_model, which
      supplies the opacity array, but the densities are passed through
      from dust_model::total_abs_length. NOTE that due to deficiencies
      in how blitz can operate on expressions, the density vector MUST
      be mapped to dimension 2, ie pass rho(tensor::j) or if an
      expression, map each of the operands to dimension 2. */
  template <typename T>
  static typename RT_total_abs_length2<T>::T_result
  total_abs_length (const blitz::ETBase<T>& rho, const T_opacity_array& kappa) {
    using namespace blitz;
    return sum(rho.unwrap()*kappa(tensor::j, tensor::i), tensor::j);
  }

  /** Computes the total absorption length (mean free path) by summing
      over all density species and puts the result in kappa_out. This
      is called from dust_model, which supplies the opacity array, but
      the densities are passed through by
      dust_model::total_abs_length. To overcome inefficiencies in the
      blitz reductions, this is calculated explicitly with
      pointers. */
  static inline void
  total_abs_length (const T_densities& rho, const T_opacity_array& kappa,
		    T_lambda& kappa_out) {
    using namespace blitz;
    assert(rho.size()==kappa.extent(firstDim));
    // kappa array minor rank is the lambda index
    assert(kappa.stride(secondDim)==1);
    T_float* restrict outptr = kappa_out.dataFirst();
    const T_float* restrict rptr = rho.dataFirst();
    const T_float* restrict kptr = kappa.dataFirst();
    const int nspecies = rho.size();
    const int nlambda = kappa_out.size();
    if(nspecies==1) {
      // if we only have one dust species, we can skip the summation loop.
      assert(kappa.stride(secondDim)==1);
      assert(kappa.extent(firstDim)==1);
#pragma ivdep
      for(int i=0; i<nlambda; i++) {
	outptr[i] = *rptr * kptr[i];
      }
    }
    else {
      // if we have more than one species, we have to sum over the
      // major rank, but the number of species should be small enough
      // that it should all fit in the cache.
      for(int i=0; i<nlambda; i++) {
	outptr[i]=0;
#pragma ivdep
	for(int j=0; j<nspecies; ++j) {
	  outptr[i]+=rptr[j] * kptr[i + j*nspecies];
	}
      }
    }
  }

  /** Traits class containing the return type definition for
      abs_length_array. */
  template <typename T> struct RT_abs_length_array {
    typedef typename blitz::BzBinaryExprResult<
      blitz::Multiply, T, T_opacity_array>::T_result T_result;
  };
  
  /** Return an array of absorption lengths (which is dtau/dl) by
      multiplying the opacity array with a density array. This is
      essentially the same as total_abs_length, except it does not sum
      over densities. NOTE that the same caveat applies about mapping
      the rho, except here it must be mapped to dimension *1*, ie pass
      rho(tensor::i). */
  template <typename T>
  static typename RT_abs_length_array<T>::T_result
  abs_length_array (const blitz::ETBase<T>& rho, const T_opacity_array& kappa)
  {
    // kappa==kappa(i,j) but without the indexing overhead.
    return rho.unwrap()*kappa;
  };

  static T_densities reference_abs_lengths (const T_densities&, const T_opacity_array&, T_biaser);
  static T_float total_reference_abs_length (const T_densities&, const T_opacity_array&, T_biaser);
  static T_lambda extract_scatterer_component (const T_opacity_array&, int);

  static void set_wavelength_prefix (T_opacity_array& kappa,
				     T_scatterer_vector& scatterers,
				     const T_lambda& lambda);
  static void set_wavelength_postfix (T_opacity_array& kappa,
				      T_scatterer_vector& scatterers,
				      const T_lambda& lambda);
};


// *** monochromatic function definitions ***

/** Return an expression template representing an array of absorption
    lengths (which is dtau/dl) obtained by multiplying the opacity
    array with a density array. */
template <typename scatterer_type> 
typename blitz::BzBinaryExprResult
<blitz::Multiply, mcrx::T_densities, mcrx::array_1>::T_result
mcrx::monochromatic_dust_model_policy<scatterer_type>::abs_length_array 
(const T_densities& rho,
 const T_opacity_array& kappa)
{
  return rho*kappa;
}

/** Calculate the total absorption length (summed over all species) by
    summing the absorption length array. */
template <typename scatterer_type> 
template<typename T>
typename mcrx::monochromatic_dust_model_policy<scatterer_type>::T_lambda 
mcrx::monochromatic_dust_model_policy<scatterer_type>::total_abs_length 
(blitz::_bz_ArrayExpr<T> kappa)
{
  return sum(kappa);
}

/** Calculate the total absorption length (summed over all species) by
    summing the absorption length array. */
template <typename scatterer_type> 
typename mcrx::monochromatic_dust_model_policy<scatterer_type>::T_lambda 
mcrx::monochromatic_dust_model_policy<scatterer_type>::total_abs_length 
(const T_opacity_array& kappa)
{
  return sum(kappa);
}

/** This function is called by scatterer::set_wavelength before
    changing the wavelengths of the scatterers themselves. */
template <typename scatterer_type>
void
mcrx::monochromatic_dust_model_policy<scatterer_type>::set_wavelength_prefix
(T_opacity_array& kappa, 
 T_scatterer_vector& scatterers, const
 T_lambda& lambda) 
{
  // assure opacity array is correctly sized
  kappa.resize(scatterers.size());
}

/** This function is called by scatterer::set_wavelength after
    changing the wavelengths of the scatterers themselves. For
    monochromatic RT, it copies the scatterer opacities into the
    dust_model opacity array. */
template <typename scatterer_type> 
void 
mcrx::monochromatic_dust_model_policy<scatterer_type>::set_wavelength_postfix
  (T_opacity_array& kappa,
   T_scatterer_vector& scatterers,
   const T_lambda& lambda)
{
  // copy scatterer opacities to opacity array
  for (int i = 0; i < scatterers.size(); ++i)
    kappa(i) = scatterers[i]->opacity();
}


/** Extracts the component belonging to scatterer i from the supplied
    opacity array. */
template <typename scatterer_type> 
typename mcrx::monochromatic_dust_model_policy<scatterer_type>::T_lambda 
mcrx::monochromatic_dust_model_policy<scatterer_type>::extract_scatterer_component
(const T_opacity_array& array, int i)
{
  return array(i);
}


// *** polychromatic function definintions ***




/** This function is called by scatterer::set_wavelength before
    changing the wavelengths of the scatterers themselves. For
    polychromatic RT, it makes sure the opacity array is the correct
    size and updates the scatterers to reference this array. */
template <typename scatterer_type> 
void 
mcrx::polychromatic_dust_model_policy<scatterer_type>::set_wavelength_prefix 
(T_opacity_array& kappa,
 T_scatterer_vector& scatterers,
 const T_lambda& lambda)
{
  // resize opacity array
  kappa.resize(scatterers.size(), lambda.size());
  // update references, while preserving the data currently in the
  // scatterer opacity array
  for (int i = 0; i < scatterers.size(); ++i) {
    if(scatterers[i]->opacity_.size() == lambda.size())
      kappa(i, blitz::Range::all()) = scatterers[i]->opacity_;
    scatterers[i]->opacity_.reference(kappa(i, blitz::Range::all()));
  }
}


/** This function is called by scatterer::set_wavelength after
    changing the wavelengths of the scatterers themselves. For
    polychromatic RT, it's a no-op. */
template <typename scatterer_type> 
void 
mcrx::polychromatic_dust_model_policy<scatterer_type>::set_wavelength_postfix
  (T_opacity_array& kappa,
   T_scatterer_vector& scatterers,
   const T_lambda& lambda)
{
  // nop
}
 

/** Extracts the component belonging to scatterer i from the supplied
    opacity array. */
template <typename scatterer_type> 
typename mcrx::polychromatic_dust_model_policy<scatterer_type>::T_lambda 
mcrx::polychromatic_dust_model_policy<scatterer_type>::extract_scatterer_component
(const T_opacity_array& array, int i)
{
  return array(i, blitz::Range::all());
}

#endif
